A Proposal to Improve the Hamiltonian Dynamics in Stan

Introduction

I propose new dynamics for the Hamiltonian Monte Carlo (HMC) algorithm implemented in the probabilistic programming language Stan. This technical report compares Stan’s version of HMC to my proposal, but does not attempt to explain Hamiltonian Monte Carlo. For a thorough introduction to HMC, consider Michael Betancourt’s A Conceptual Introduction to Hamiltonian Monte Carlo.

After describing the differences between these two versions of HMC, a simulation study is carried out. XXX Describe the simulation study: what are we comparing and why. Many of the models used in this study can be found in the GitHub repository stan-dev/performance-tests-cmdstan and specifically in the folder stat_comp_benchmarks therein. A few variations on Gaussian and studten-t distributions are also considered.

Proposed Dynamics

Let \(\pi(q) \propto \exp{(-U(q))}\) be the target distribution of interest for parameters \(q \in \mathbb{R}^{d}\). As of 2020-04-22, the Stan GitHub branch develop’s version of HMC could be written as a Hamiltonian

\[ H_{dv}(q, p) = U(q) + \frac{1}{2} p^T M^{-1} p \]

where the subscript \(dv\) stands for develop, \(M^{-1} \approx \mathbb{E}_{\pi} (q - \mu)(q - \mu)^T\). Under this scheme, \(p \sim N(0, M)\).

The work of Ma, Chen, and Fox (2015) in A Complete Recipe for Stochastic Gradient MCMC suggests a richer framework than shown below. They also offer a stochastic version of this, but I’m ignoring all of that. In an effort to simplify Ma’s notation, I’ve removed some terms from the equations below that will drop out eventually anyway of what I’m proposing.

Let \(z = (q, p)\). Ma et al.’s framework dictates dynamics as

\[\mathrm{d}z = -Q(z) \nabla H(z)\]

where \(Q\) is any skew symmetric matrix. Some choices of \(Q\) will lead to faster convergence than others.

Branch develop’s version of HMC, in the notation of this framework, goes like this. Draw \(p \sim N(0, M)\), necessitating a decomposition \(LL^{T} = M^{-1}\) and a solve(). This was chosen by design because \(M^{-1}\) itself is needed in the next step, and it’s expensive to keep both \(L\) and \(M^{-1}\) in memory. With \(p\) in hand and

\[Q(z) = \begin{bmatrix} 0 && -I \\ I && 0 \end{bmatrix}\]

update \(z\) according to the dynamcis

\[ \mathrm{d}z = -Q(z) \nabla U(q) = \begin{bmatrix} 0 && -I \\ I && 0 \end{bmatrix} \begin{bmatrix} \nabla U(q) \\ M^{-1}p \end{bmatrix} = \begin{bmatrix} M^{-1}p \\ -\nabla U(q) \end{bmatrix}.\]

Stan appproximates these dynamics using the half-step leapfrog integrator and this strategy will apply throughout this discussion.

The proposed version of HMC draws \(p_t \sim N(0, I)\), thus specifying a slightly modified Hamiltonian,

$ \(H_{pr}(z) = U(q) + \frac{1}{2}p^Tp.\)$

Take

\[ Q(z) = \begin{bmatrix} 0 && -L^T \\ L && 0 \end{bmatrix} \]

where \(LL^T = M^{-1}\). Since \(M^{-1}\) is only updated at the end of each adaption window, \(L\) is fixed throughout sampling. There is no need for \(M^{-1}\) itself. The proposed dynamics then follow

\[ \mathrm{d}z = \begin{bmatrix} L^{T}p \\ -L \nabla U(q) \end{bmatrix}.\]

Overall, the proposed dynamics should be faster as there are fewer decompositions and there is no matrix inverse via a call to solve(). While \(L \cdot \nabla U(q)\) is calculated twice within the leapfrog integrator, once for each half-step, the matrix vector multiplication should be more efficient in Eigen since \(L\) is diagonal¹.

Numerical Study

Below is a comparison of the two version of HMC dynamics presented above, for both XXX change this metrics dense_e and diag_e across a range of models. All models were fit with CmdStan v2.22.1 on a MacBook Pro running macOS 10.14.6, which has a 2GHz Intel Core i5 processor and 16GB 1867 MHz LPDDR3 memory. Each model was fit with 2000 warmup iterations, 10000 iterations beyond warmup, and 4 chains. All other Stan defaults were left alone.

XXX Repeat purpose of study and state where proposed code can be found.

Model: anisotropic_gauss_c25_100d

Runtime

The model agc25 has 101 parameters. The table below displays the run time of this model in seconds.

Branch	seconds
develop	116
proposal	51

ESS

Rhat

Summary Statistics

Estimated parameters.

Parameter Monte Carlo standard errors.

metric	MCSE_Q2.5	MCSE_Q25	MCSE_Q50	MCSE_Q75	MCSE_Q97.5	MCSE_SD
lp__
dense_e	0.017	0.005	0.005	0.007	0.014	0.005
dense_e	0.013	0.007	0.006	0.007	0.017	0.005
x[1]
dense_e	0.018	0.006	0.006	0.006	0.014	0.005
dense_e	0.016	0.006	0.005	0.007	0.012	0.005
x[10]
dense_e	0.013	0.007	0.005	0.007	0.011	0.005
dense_e	0.014	0.007	0.005	0.007	0.015	0.005
x[100]
dense_e	0.015	0.006	0.005	0.007	0.012	0.005
dense_e	0.017	0.006	0.005	0.006	0.013	0.005
x[11]
dense_e	0.012	0.007	0.005	0.007	0.013	0.005
dense_e	0.013	0.006	0.005	0.007	0.015	0.005
x[12]
dense_e	0.013	0.006	0.005	0.006	0.016	0.005
dense_e	0.013	0.006	0.006	0.007	0.011	0.005
x[13]
dense_e	0.015	0.007	0.005	0.007	0.012	0.005
dense_e	0.017	0.007	0.005	0.008	0.016	0.005
x[14]
dense_e	0.012	0.007	0.005	0.008	0.016	0.005
dense_e	0.017	0.006	0.005	0.006	0.015	0.005
x[15]
dense_e	0.015	0.007	0.006	0.007	0.016	0.005
dense_e	0.011	0.007	0.006	0.007	0.014	0.005
x[16]
dense_e	0.014	0.008	0.004	0.006	0.018	0.005
dense_e	0.013	0.007	0.005	0.006	0.014	0.005
x[17]
dense_e	0.014	0.007	0.005	0.007	0.014	0.005
dense_e	0.014	0.006	0.006	0.007	0.013	0.005
x[18]
dense_e	0.016	0.007	0.005	0.007	0.015	0.005
dense_e	0.016	0.006	0.005	0.006	0.013	0.005
x[19]
dense_e	0.015	0.006	0.006	0.007	0.014	0.005
dense_e	0.016	0.006	0.005	0.007	0.012	0.005
x[2]
dense_e	0.015	0.007	0.005	0.006	0.014	0.005
dense_e	0.014	0.005	0.005	0.007	0.013	0.005
x[20]
dense_e	0.015	0.006	0.006	0.007	0.016	0.005
dense_e	0.016	0.007	0.005	0.007	0.014	0.005
x[21]
dense_e	0.013	0.007	0.005	0.007	0.014	0.005
dense_e	0.018	0.007	0.005	0.007	0.015	0.005
x[22]
dense_e	0.016	0.007	0.005	0.006	0.017	0.005
dense_e	0.018	0.006	0.006	0.007	0.012	0.005
x[23]
dense_e	0.014	0.006	0.006	0.007	0.016	0.005
dense_e	0.015	0.007	0.005	0.007	0.016	0.005
x[24]
dense_e	0.013	0.008	0.005	0.007	0.014	0.005
dense_e	0.016	0.007	0.006	0.007	0.018	0.005
x[25]
dense_e	0.012	0.007	0.005	0.006	0.015	0.005
dense_e	0.017	0.006	0.005	0.006	0.013	0.005
x[26]
dense_e	0.016	0.007	0.006	0.006	0.012	0.005
dense_e	0.014	0.007	0.005	0.006	0.014	0.005
x[27]
dense_e	0.014	0.008	0.006	0.007	0.016	0.005
dense_e	0.015	0.006	0.006	0.007	0.012	0.005
x[28]
dense_e	0.017	0.006	0.005	0.007	0.013	0.005
dense_e	0.013	0.007	0.005	0.007	0.015	0.005
x[29]
dense_e	0.013	0.007	0.005	0.006	0.015	0.005
dense_e	0.014	0.007	0.006	0.007	0.015	0.005
x[3]
dense_e	0.015	0.007	0.006	0.006	0.013	0.005
dense_e	0.015	0.007	0.006	0.007	0.015	0.005
x[30]
dense_e	0.013	0.006	0.005	0.006	0.015	0.005
dense_e	0.015	0.008	0.005	0.007	0.013	0.005
x[31]
dense_e	0.016	0.007	0.006	0.006	0.013	0.005
dense_e	0.015	0.008	0.005	0.006	0.014	0.005
x[32]
dense_e	0.015	0.006	0.005	0.006	0.015	0.005
dense_e	0.016	0.007	0.006	0.006	0.014	0.005
x[33]
dense_e	0.012	0.006	0.005	0.007	0.012	0.005
dense_e	0.015	0.007	0.006	0.006	0.018	0.005
x[34]
dense_e	0.012	0.007	0.005	0.007	0.015	0.005
dense_e	0.014	0.007	0.006	0.007	0.014	0.005
x[35]
dense_e	0.015	0.006	0.006	0.007	0.013	0.005
dense_e	0.014	0.007	0.006	0.006	0.016	0.005
x[36]
dense_e	0.015	0.007	0.005	0.007	0.016	0.005
dense_e	0.012	0.007	0.006	0.006	0.016	0.005
x[37]
dense_e	0.013	0.006	0.005	0.007	0.015	0.005
dense_e	0.015	0.006	0.004	0.006	0.019	0.005
x[38]
dense_e	0.016	0.007	0.005	0.006	0.013	0.005
dense_e	0.014	0.006	0.006	0.008	0.015	0.005
x[39]
dense_e	0.017	0.007	0.006	0.008	0.016	0.005
dense_e	0.016	0.006	0.005	0.007	0.014	0.005
x[4]
dense_e	0.014	0.007	0.006	0.007	0.013	0.005
dense_e	0.014	0.007	0.005	0.007	0.016	0.005
x[40]
dense_e	0.014	0.006	0.005	0.007	0.015	0.005
dense_e	0.014	0.007	0.005	0.007	0.018	0.005
x[41]
dense_e	0.014	0.007	0.005	0.007	0.014	0.005
dense_e	0.015	0.007	0.005	0.007	0.011	0.005
x[42]
dense_e	0.018	0.006	0.005	0.006	0.014	0.005
dense_e	0.016	0.007	0.005	0.007	0.015	0.005
x[43]
dense_e	0.016	0.006	0.005	0.007	0.017	0.005
dense_e	0.018	0.008	0.006	0.007	0.013	0.005
x[44]
dense_e	0.015	0.008	0.005	0.006	0.015	0.005
dense_e	0.015	0.008	0.006	0.006	0.015	0.005
x[45]
dense_e	0.014	0.006	0.006	0.007	0.014	0.005
dense_e	0.017	0.007	0.005	0.007	0.011	0.005
x[46]
dense_e	0.016	0.007	0.006	0.006	0.015	0.005
dense_e	0.016	0.007	0.005	0.006	0.014	0.005
x[47]
dense_e	0.014	0.007	0.005	0.007	0.013	0.005
dense_e	0.014	0.006	0.005	0.006	0.011	0.005
x[48]
dense_e	0.014	0.007	0.006	0.006	0.012	0.005
dense_e	0.015	0.008	0.005	0.006	0.015	0.005
x[49]
dense_e	0.014	0.007	0.006	0.007	0.015	0.005
dense_e	0.015	0.007	0.005	0.006	0.019	0.005
x[5]
dense_e	0.016	0.007	0.006	0.007	0.015	0.005
dense_e	0.012	0.007	0.006	0.007	0.013	0.005
x[50]
dense_e	0.016	0.006	0.006	0.007	0.013	0.005
dense_e	0.018	0.007	0.005	0.006	0.015	0.005
x[51]
dense_e	0.014	0.006	0.006	0.007	0.016	0.005
dense_e	0.015	0.007	0.005	0.007	0.016	0.005
x[52]
dense_e	0.015	0.006	0.005	0.006	0.012	0.005
dense_e	0.013	0.006	0.005	0.007	0.014	0.005
x[53]
dense_e	0.015	0.007	0.005	0.006	0.015	0.005
dense_e	0.014	0.007	0.005	0.007	0.016	0.005
x[54]
dense_e	0.014	0.006	0.006	0.007	0.012	0.005
dense_e	0.014	0.008	0.005	0.007	0.015	0.005
x[55]
dense_e	0.013	0.006	0.005	0.008	0.016	0.005
dense_e	0.013	0.006	0.006	0.007	0.016	0.005
x[56]
dense_e	0.013	0.006	0.006	0.006	0.012	0.005
dense_e	0.018	0.007	0.005	0.006	0.015	0.005
x[57]
dense_e	0.013	0.008	0.005	0.007	0.014	0.005
dense_e	0.013	0.006	0.006	0.006	0.014	0.005
x[58]
dense_e	0.016	0.007	0.006	0.007	0.016	0.005
dense_e	0.016	0.007	0.006	0.006	0.017	0.005
x[59]
dense_e	0.015	0.007	0.006	0.007	0.014	0.005
dense_e	0.016	0.007	0.006	0.007	0.013	0.005
x[6]
dense_e	0.014	0.008	0.005	0.006	0.014	0.005
dense_e	0.012	0.007	0.006	0.008	0.013	0.005
x[60]
dense_e	0.016	0.007	0.005	0.006	0.015	0.005
dense_e	0.016	0.007	0.005	0.007	0.012	0.005
x[61]
dense_e	0.017	0.006	0.004	0.006	0.016	0.005
dense_e	0.015	0.007	0.006	0.006	0.015	0.005
x[62]
dense_e	0.014	0.007	0.005	0.006	0.018	0.005
dense_e	0.015	0.006	0.006	0.007	0.015	0.005
x[63]
dense_e	0.014	0.006	0.005	0.006	0.013	0.005
dense_e	0.016	0.007	0.006	0.007	0.016	0.005
x[64]
dense_e	0.014	0.007	0.005	0.007	0.016	0.005
dense_e	0.015	0.007	0.005	0.007	0.012	0.005
x[65]
dense_e	0.012	0.007	0.005	0.007	0.016	0.005
dense_e	0.015	0.008	0.005	0.006	0.016	0.005
x[66]
dense_e	0.014	0.007	0.005	0.007	0.018	0.005
dense_e	0.013	0.006	0.005	0.007	0.014	0.005
x[67]
dense_e	0.017	0.007	0.005	0.007	0.014	0.005
dense_e	0.012	0.006	0.006	0.006	0.015	0.005
x[68]
dense_e	0.018	0.006	0.005	0.007	0.014	0.005
dense_e	0.017	0.006	0.006	0.007	0.015	0.005
x[69]
dense_e	0.013	0.006	0.005	0.007	0.015	0.005
dense_e	0.015	0.007	0.005	0.007	0.013	0.005
x[7]
dense_e	0.012	0.007	0.005	0.007	0.012	0.005
dense_e	0.017	0.006	0.006	0.006	0.016	0.005
x[70]
dense_e	0.014	0.006	0.005	0.007	0.016	0.005
dense_e	0.014	0.006	0.006	0.007	0.014	0.005
x[71]
dense_e	0.013	0.006	0.005	0.007	0.015	0.005
dense_e	0.015	0.007	0.006	0.006	0.017	0.005
x[72]
dense_e	0.019	0.006	0.005	0.007	0.014	0.005
dense_e	0.014	0.007	0.006	0.007	0.014	0.005
x[73]
dense_e	0.018	0.007	0.005	0.007	0.014	0.005
dense_e	0.012	0.006	0.005	0.007	0.014	0.005
x[74]
dense_e	0.016	0.006	0.006	0.007	0.015	0.005
dense_e	0.018	0.008	0.006	0.007	0.014	0.005
x[75]
dense_e	0.014	0.007	0.005	0.007	0.014	0.005
dense_e	0.013	0.006	0.006	0.007	0.015	0.005
x[76]
dense_e	0.013	0.007	0.006	0.007	0.012	0.005
dense_e	0.017	0.007	0.006	0.006	0.013	0.005
x[77]
dense_e	0.016	0.006	0.006	0.007	0.014	0.005
dense_e	0.016	0.007	0.005	0.006	0.014	0.005
x[78]
dense_e	0.014	0.006	0.005	0.007	0.017	0.005
dense_e	0.014	0.006	0.005	0.008	0.013	0.005
x[79]
dense_e	0.014	0.007	0.005	0.007	0.013	0.005
dense_e	0.016	0.007	0.005	0.007	0.017	0.005
x[8]
dense_e	0.014	0.007	0.005	0.007	0.016	0.005
dense_e	0.014	0.006	0.006	0.008	0.014	0.005
x[80]
dense_e	0.015	0.007	0.006	0.007	0.016	0.005
dense_e	0.014	0.007	0.006	0.006	0.015	0.005
x[81]
dense_e	0.017	0.006	0.005	0.007	0.011	0.005
dense_e	0.016	0.006	0.006	0.006	0.016	0.005
x[82]
dense_e	0.016	0.007	0.006	0.007	0.017	0.005
dense_e	0.016	0.006	0.006	0.007	0.016	0.005
x[83]
dense_e	0.019	0.008	0.005	0.007	0.012	0.005
dense_e	0.014	0.007	0.006	0.008	0.013	0.005
x[84]
dense_e	0.016	0.006	0.006	0.006	0.013	0.005
dense_e	0.013	0.007	0.006	0.006	0.017	0.005
x[85]
dense_e	0.012	0.007	0.005	0.006	0.013	0.005
dense_e	0.013	0.007	0.006	0.007	0.015	0.005
x[86]
dense_e	0.014	0.006	0.005	0.007	0.015	0.005
dense_e	0.017	0.007	0.005	0.006	0.014	0.005
x[87]
dense_e	0.013	0.007	0.005	0.006	0.015	0.005
dense_e	0.014	0.006	0.005	0.006	0.015	0.005
x[88]
dense_e	0.016	0.007	0.005	0.007	0.014	0.005
dense_e	0.014	0.007	0.006	0.007	0.018	0.004
x[89]
dense_e	0.012	0.006	0.006	0.006	0.014	0.005
dense_e	0.012	0.007	0.006	0.007	0.016	0.005
x[9]
dense_e	0.012	0.007	0.005	0.006	0.011	0.005
dense_e	0.013	0.008	0.006	0.008	0.012	0.005
x[90]
dense_e	0.015	0.007	0.005	0.006	0.014	0.005
dense_e	0.016	0.007	0.006	0.007	0.015	0.005
x[91]
dense_e	0.017	0.007	0.005	0.006	0.012	0.005
dense_e	0.017	0.007	0.006	0.007	0.015	0.005
x[92]
dense_e	0.015	0.007	0.005	0.006	0.015	0.005
dense_e	0.017	0.007	0.006	0.007	0.013	0.005
x[93]
dense_e	0.015	0.007	0.005	0.006	0.015	0.005
dense_e	0.013	0.008	0.007	0.007	0.013	0.005
x[94]
dense_e	0.015	0.008	0.006	0.006	0.019	0.005
dense_e	0.016	0.008	0.007	0.008	0.016	0.005
x[95]
dense_e	0.014	0.006	0.005	0.006	0.017	0.005
dense_e	0.012	0.007	0.007	0.007	0.014	0.005
x[96]
dense_e	0.013	0.007	0.006	0.008	0.014	0.005
dense_e	0.014	0.007	0.007	0.008	0.014	0.005
x[97]
dense_e	0.015	0.007	0.005	0.006	0.016	0.005
dense_e	0.017	0.007	0.007	0.008	0.014	0.005
x[98]
dense_e	0.014	0.007	0.005	0.007	0.013	0.005
dense_e	0.018	0.007	0.007	0.008	0.015	0.005
x[99]
dense_e	0.133	0.066	0.071	0.064	0.103	0.042
dense_e	0.129	0.074	0.062	0.059	0.079	0.041

Model: arK

Runtime

The model arK has 8 parameters. The table below displays the run time of this model in seconds.

Branch	seconds
develop	6
proposal	11

ESS

Rhat

Summary Statistics

Estimated parameters.

Parameter Monte Carlo standard errors.

metric	MCSE_Q2.5	MCSE_Q25	MCSE_Q50	MCSE_Q75	MCSE_Q97.5	MCSE_SD
alpha
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
beta[1]
dense_e	0.001	0.000	0.000	0.001	0.001	0.000
dense_e	0.001	0.001	0.001	0.001	0.001	0.000
beta[2]
dense_e	0.001	0.001	0.000	0.001	0.001	0.000
dense_e	0.001	0.001	0.001	0.001	0.001	0.000
beta[3]
dense_e	0.001	0.001	0.000	0.001	0.001	0.000
dense_e	0.002	0.001	0.001	0.001	0.002	0.001
beta[4]
dense_e	0.001	0.001	0.000	0.001	0.001	0.000
dense_e	0.002	0.001	0.001	0.002	0.002	0.001
beta[5]
dense_e	0.001	0.001	0.000	0.000	0.001	0.000
dense_e	0.004	0.002	0.002	0.002	0.003	0.002
lp__
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
sigma
dense_e	0.047	0.021	0.017	0.012	0.010	0.010
dense_e	0.078	0.031	0.021	0.016	0.011	0.015

Model: arma

Runtime

The model arma has 5 parameters. The table below displays the run time of this model in seconds.

Branch	seconds
develop	3
proposal	3

ESS

Rhat

Summary Statistics

Estimated parameters.

Parameter Monte Carlo standard errors.

metric	MCSE_Q2.5	MCSE_Q25	MCSE_Q50	MCSE_Q75	MCSE_Q97.5	MCSE_SD
lp__
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
mu
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
phi
dense_e	0.001	0.000	0.000	0.000	0.001	0.000
dense_e	0.001	0.000	0.000	0.000	0.001	0.000
sigma
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
dense_e	0.000	0.000	0.000	0.000	0.000	0.000
theta
dense_e	0.040	0.015	0.011	0.008	0.005	0.007
dense_e	0.041	0.018	0.011	0.008	0.004	0.007

References

Betancourt, Michael. 2017. “A Conceptual Introduction to Hamiltonian Monte Carlo.” arXiv Preprint arXiv:1701.02434.

Leimkuhler, Benedict, Charles Matthews, and Jonathan Weare. 2018. “Ensemble Preconditioning for Markov Chain Monte Carlo Simulation.” Statistics and Computing 28 (2). Springer: 277–90.

Ma, Yi-An, Tianqi Chen, and Emily Fox. 2015. “A Complete Recipe for Stochastic Gradient Mcmc.” In Advances in Neural Information Processing Systems, 2917–25.

Team, Stan Development. 2018. “The Stan Core Library.” http://mc-stan.org.

---

1. Thanks to Matija Rezar and Ben Bales for pointing this out.↩